import numpy as np
from scipy import signal
from scipy.fft import fft, fftfreq
import matplotlib.pyplot as plt

def calculate_power_high_precision(i_signal, q_signal, fs, signal_freq, n_fft_mult=4, impedance=50.0):
    """
    使用高精度方法计算I/Q信号的功率
    
    参数:
    i_signal (np.ndarray): I路信号
    q_signal (np.ndarray): Q路信号
    fs (float): 采样率 (Hz)
    signal_freq (float): 目标信号频率 (Hz)
    n_fft_mult (int): FFT补零倍数，提高频率分辨率
    impedance (float): 系统阻抗 (Ω)，默认50Ω
    
    返回:
    float: 信号功率 (dBm)
    """
    # 检查输入信号长度一致性
    if len(i_signal) != len(q_signal):
        raise ValueError("I信号和Q信号长度必须相同")
    
    N = len(i_signal)
    # 构建复数信号
    complex_signal = i_signal + 1j * q_signal
    # 去除直流分量
    complex_signal = complex_signal - np.mean(complex_signal)

    # 生成平顶窗并计算补偿系数
    window = signal.windows.flattop(N)
    coherent_gain = np.mean(window)  # 窗函数的相干增益

    # 加窗并进行FFT（补零以提高频率分辨率）
    n_fft = N * n_fft_mult
    spectrum = fft(complex_signal * window, n=n_fft) / (N * coherent_gain)
    spectrum_abs = np.abs(spectrum)  # 双边谱幅度
    
    # 单边谱转换（只对实数信号需要乘以2，复数信号不需要）
    # spectrum_abs = spectrum_abs[:n_fft // 2] * 2  # 注释掉这行！
    
    freq_axis = fftfreq(n_fft, 1 / fs)  # 双边频率轴

    # 找到目标频率附近的峰值（双边谱）
    peak_idx = np.argmin(np.abs(freq_axis - signal_freq))
    
    # 使用三阶多项式插值提高峰值估计精度
    idx_range = 3  # 插值点范围
    idx_low = max(0, peak_idx - idx_range)
    idx_high = min(len(freq_axis), peak_idx + idx_range + 1)
    
    # 多项式拟合
    poly = np.polyfit(freq_axis[idx_low:idx_high], spectrum_abs[idx_low:idx_high], 3)
    true_peak = np.polyval(poly, signal_freq)  # 在目标频率处的插值结果

    # 绘制频谱图（可选）
    plt.figure(figsize=(10, 6))
    plt.plot(freq_axis, 20 * np.log10(spectrum_abs))
    plt.axvline(signal_freq, color='r', linestyle='--', label=f'目标频率: {signal_freq} Hz')
    plt.xlabel("Frequency (Hz)")
    plt.ylabel("Amplitude (dB)")
    plt.title("复数信号频谱")
    plt.grid(True)
    plt.legend()
    plt.show()

    # 计算功率
    # 对于复数信号，功率 = 幅度²/阻抗，转换为dBm
    power_watts = (true_peak ** 2) / impedance
    power_dbm = 10 * np.log10(power_watts * 1000)  # 转换为dBm

    return power_dbm


# 测试代码
if __name__ == "__main__":
    # 生成测试I/Q信号（50.3Hz的复数正弦波，幅度1.0V）
    fs = 1000  # 采样率 (Hz)
    f0 = 50.3  # 信号频率 (Hz)
    N = 4096   # 采样点数
    t = np.arange(N) / fs
    
    # 生成理想I/Q信号（幅度1.0V，相位正交）
    amplitude = 1.0
    i_signal = amplitude * np.cos(2 * np.pi * f0 * t)
    q_signal = amplitude * np.sin(2 * np.pi * f0 * t)

    # 计算功率
    power = calculate_power_high_precision(i_signal, q_signal, fs, f0, n_fft_mult=4)
    
    # 理论功率计算
    # 对于复数信号，理论功率 = 幅度²/阻抗
    theoretical_power = (amplitude ** 2) / 50.0  # 假设50Ω阻抗
    theoretical_dbm = 10 * np.log10(theoretical_power * 1000)
    
    # 输出结果
    print(f"计算功率: {power:.3f} dBm")
    print(f"理论功率: {theoretical_dbm:.3f} dBm")
    print(f"绝对误差: {power - theoretical_dbm:.3f} dB")    